# __author__ = 'heyin'
# __date__ = '2018/11/12 15:38'
# 朴素贝叶斯对文本进行分类
from sklearn.datasets import fetch_20newsgroups
from sklearn.decomposition import PCA
from sklearn.metrics import classification_report
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import TfidfVectorizer, CountVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.feature_selection import VarianceThreshold


def news_bayes():
    news = fetch_20newsgroups(subset='all')
    # 划分数据集
    x_train, x_test, y_train, y_test = train_test_split(news.data, news.target, test_size=0.25)
    # 文本特征抽取，使用tf-idf
    tf = TfidfVectorizer()
    # 为了节省内存，这些默认都是稀疏矩阵，如果使用toarray，会提示内存错误
    x_train = tf.fit_transform(x_train)
    x_test = tf.transform(x_test)
    print(x_train.shape)  # (14134, 153483)
    # 因为预测结果表现上看，在训练集一直优于测试集大于7-8个百分点，考虑过拟合
    # 发现有15w多的特征，考虑降维，pca不能使用sparse矩阵，因此测试失败
    # pca = PCA(n_components=0.9)
    # x_train = pca.fit_transform(x_train)
    # x_test = pca.transform(x_test)

    # 过滤式降维，即特征选择没有效果
    # vt = VarianceThreshold(threshold=0)
    # x_train = vt.fit_transform(x_train)
    # x_test = vt.transform(x_test)

    print(x_train.shape)

    nb = MultinomialNB(alpha=1)  # 没有超参数需要调优
    nb.fit(x_train, y_train)
    y_pred = nb.predict(x_test)
    print('score:', nb.score(x_test, y_test))
    print('训练集score:', nb.score(x_train, y_train))
    # print('预测结果精度汇总：', classification_report(y_test, y_pred))


def cv_bayes():
    news = fetch_20newsgroups(subset='all')
    # 划分数据集
    x_train, x_test, y_train, y_test = train_test_split(news.data, news.target, test_size=0.25)
    # 文本特征抽取，使用tf-idf
    cv = CountVectorizer()
    # 为了节省内存，这些默认都是稀疏矩阵，如果使用toarray，会提示内存错误
    x_train = cv.fit_transform(x_train)
    x_test = cv.transform(x_test)

    nb = MultinomialNB(alpha=1)  # 没有超参数需要调优
    nb.fit(x_train, y_train)
    y_pred = nb.predict(x_test)
    print('测试集score:', nb.score(x_test, y_test))
    print('训练集score:', nb.score(x_train, y_train))
    # print('预测结果精度汇总：', classification_report(y_test, y_pred))


if __name__ == '__main__':
    news_bayes()
    cv_bayes()
